{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f431da81-4d94-4fd1-8a51-952ef824cb46",
   "metadata": {},
   "source": [
    "# IMDB Movie Review Dataset\n",
    "\n",
    "This is a library to download and parse the [IMDB's Large Movie Review Dataset](http://ai.stanford.edu/~amaas/data/sentiment/) dataset and a demo of a transformer based model. The dataset has 25K training, and 25K test dataset, plus 50K unlabeled examples.\n",
    "\n",
    "It's inspired on [Keras' Text classification with Transformer](https://keras.io/examples/nlp/text_classification_with_transformer/) demo."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "392f19ad-caa7-4d7d-8078-cbca680f2be2",
   "metadata": {},
   "source": [
    "## Environment Set Up\n",
    "\n",
    "Let's set up `go.mod` to use the local copy of GoMLX, so it can be developed jointly the dataset code with the model. That's often how data pre-processing and model code is developed together with experimentation.\n",
    "\n",
    "If you are not changing code, feel free to simply skip this cell. Or if you used a different directory for you projects, change it below.\n",
    "\n",
    "Notice the directory `${HOME}/Projects/gomlx` is where the GoMLX code is copied by default in [its Docker](https://hub.docker.com/repository/docker/janpfeifer/gomlx_jupyterlab/general)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "32600e05-745d-4c38-999a-3b96ca1502cc",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t- Added replace rule for module \"github.com/gomlx/gomlx\" to local directory \"/home/janpf/Projects/gomlx\".\n",
      "\t- Added replace rule for module \"github.com/gomlx/gopjrt\" to local directory \"/home/janpf/Projects/gopjrt\".\n"
     ]
    }
   ],
   "source": [
    "!*rm -f go.work && go work init && go work use . \"${HOME}/Projects/gomlx\" \"${HOME}/Projects/gopjrt\"\n",
    "%goworkfix"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "443e9e56-ef12-4d04-b7bb-99478fa2b662",
   "metadata": {},
   "source": [
    "## Data Preparation\n",
    "\n",
    "### Downloading data files\n",
    "\n",
    "To download, uncompress and untar to the local directory, simply do the following. Notice if it's already downloaded in the given `--data` directory, it returns immediately."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e98ca191-0d4b-4e7c-ab24-ded6991b9671",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> Loading previously generated preprocessed binary file.\n",
      "Loaded data from \"aclImdb.bin\": 100000 examples, 141088 unique tokens, 23727054 tokens in total.\n"
     ]
    }
   ],
   "source": [
    "import (\n",
    "    \"github.com/gomlx/gomlx/examples/imdb\"\n",
    "    \"github.com/gomlx/gomlx/pkg/support/fsutil\"\n",
    "    \"github.com/janpfeifer/must\"\n",
    "\n",
    "    _ \"github.com/gomlx/gomlx/backends/default\"\n",
    ")\n",
    "\n",
    "var (\n",
    "\tflagDataDir    = flag.String(\"data\", \"~/tmp/imdb\", \"Directory to cache downloaded and generated dataset files.\")\n",
    "\tflagEval       = flag.Bool(\"eval\", true, \"Whether to evaluate the model on the validation data in the end.\")\n",
    "\tflagVerbosity  = flag.Int(\"verbosity\", 1, \"Level of verbosity, the higher the more verbose.\")\n",
    "\tflagCheckpoint = flag.String(\"checkpoint\", \"\", \"Directory save and load checkpoints from. If left empty, no checkpoints are created.\")\n",
    ")\n",
    "\n",
    "func AssertDownloaded() {\n",
    "    *flagDataDir = must.M1(fsutil.ReplaceTildeInDir(*flagDataDir))\n",
    "    if !fsutil.MustFileExists(*flagDataDir) {\n",
    "        must.M(os.MkdirAll(*flagDataDir, 0777))\n",
    "    }\n",
    "    must.M(imdb.Download(*flagDataDir))\n",
    "}\n",
    "\n",
    "%%\n",
    "AssertDownloaded()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41a415e2-8a29-4aeb-8809-b8d6b7ede5ad",
   "metadata": {},
   "source": [
    "### Sampling some examples\n",
    "\n",
    "It creates a small dataset and print out some random examples.\n",
    "\n",
    "It also defines the `DType`, used for all internal representations of the model, and the flag `--max_len` that defines the maximum number of tokens used per observation. This will beused in the modeling later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f2bf2aab-7d50-4110-a8e3-bbee08063521",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> Loading previously generated preprocessed binary file.\n",
      "Loaded data from \"aclImdb.bin\": 100000 examples, 141088 unique tokens, 23727054 tokens in total.\n",
      "┌────────────────────────────────────────────────────────────┐\n",
      "│                                                            │\n",
      "│    [Sample 0 - label 0]                                    │\n",
      "│    poorly filmed just so you don t think i can t find      │\n",
      "│    anything good here on the plus side anne jeffreys is    │\n",
      "│    very sexy in her all too brief parts of this film       │\n",
      "│    not sure if it is actually her singing or someone       │\n",
      "│    else but whoever it was had a very pretty voice ms      │\n",
      "│    jeffreys also had a couple of nice acting moments       │\n",
      "│    the script needed either a lot more of her or to        │\n",
      "│    remove her character altogether as it was her nice      │\n",
      "│    few moments weren t enough to help the film lastly      │\n",
      "│    there is mr scott he looks fantastic in this film       │\n",
      "│    and is the no nonsense lawman out to set things         │\n",
      "│    right some folks complain that his characters prior     │\n",
      "│    to 1950 were too goody goody perfect but that s         │\n",
      "│    never bothered me at all i ll take him goody goody      │\n",
      "│    pre 1950 or gritty and violent post 1950 either way     │\n",
      "│    randolph scott was a real western hero it saddens me    │\n",
      "│    to have to say it but i would have to recommend         │\n",
      "│    passing this film by unless you are a die hard fan      │\n",
      "│    there are so many better scott films out there that     │\n",
      "│    this one won t be missed                                │\n",
      "│                                                            │\n",
      "│                                                            │\n",
      "└────────────────────────────────────────────────────────────┘\n",
      "┌────────────────────────────────────────────────────────────┐\n",
      "│                                                            │\n",
      "│    [Sample 1 - label 1]                                    │\n",
      "│    <START> after reading about this documentary i          │\n",
      "│    rented it and watched it with my teenage children it    │\n",
      "│    was amazingly well balanced showing each side s         │\n",
      "│    perspective and leaving many questions unanswered       │\n",
      "│    this is as it should be i don t watch a documentary     │\n",
      "│    to be told what to think i watch it to learn and to     │\n",
      "│    draw my own conclusions afterward we took a trip to     │\n",
      "│    waco and visited the branch davidians at the site of    │\n",
      "│    the conflagration this was a potent lesson in seeing    │\n",
      "│    for yourself and it drove home the basic underlying     │\n",
      "│    honesty behind waco the rules of engagement if you      │\n",
      "│    just open your eyes and look the facts can speak for    │\n",
      "│    themselves                                              │\n",
      "│                                                            │\n",
      "│                                                            │\n",
      "└────────────────────────────────────────────────────────────┘\n",
      "┌────────────────────────────────────────────────────────────┐\n",
      "│                                                            │\n",
      "│    [Sample 2 - label 0]                                    │\n",
      "│    but they do have publicity agent peter potter           │\n",
      "│    bernard bresslaw and palace hotel owner connie joan     │\n",
      "│    sims soon enough the young beautiful wannabe models     │\n",
      "│    show up including hope springs barbara windsor paula    │\n",
      "│    perkins valerie leon dawn brakes goldfinger s           │\n",
      "│    margaret nolan debra sally geeson and ida downs         │\n",
      "│    eastenders wendy richard when the girls have cat        │\n",
      "│    fights it does draw away regular residents but after    │\n",
      "│    quite a while of some plodding not that funny           │\n",
      "│    innuendos and william jack douglas having over          │\n",
      "│    active twitches it does finally reach the               │\n",
      "│    competition and it s just afterwards i couldn t be      │\n",
      "│    bothered also starring patsy rowlands as mildred        │\n",
      "│    bumble peter butterworth as admiral joan hickson as     │\n",
      "│    mrs dukes david lodge as police inspector angela        │\n",
      "│    grant as miss bangor arnold ridley as councillor        │\n",
      "│    pratt robin askwith as larry patricia franklin as       │\n",
      "│    rosemary jimmy logan as cecil and dad s army s bill     │\n",
      "│    pertwee as fire brigade chief charles hawtrey had       │\n",
      "│    obviously quit the carry ons but where s kenneth        │\n",
      "│    williams i suppose seeing babs and young beautiful      │\n",
      "│    looking pauline fowler in bikinis but for comedy        │\n",
      "│    value this fails miserably and the overuse of the       │\n",
      "│    swanny whistle just gets on your nerves pretty poor     │\n",
      "│                                                            │\n",
      "│                                                            │\n",
      "└────────────────────────────────────────────────────────────┘\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import \"github.com/gomlx/gomlx/examples/imdb\"\n",
    "\n",
    "%%\n",
    "AssertDownloaded()\n",
    "imdb.PrintSample(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3cb8ca01-2c8f-4cd8-a644-850e928f64e1",
   "metadata": {},
   "source": [
    "## Training\n",
    "\n",
    "We will create 3 different types of models for this demo: **Bag of Words** (**\"bow\"**), **Convolutionals** (**\"cnn\"**) and **Transformers** (**\"transformer\"**).\n",
    "\n",
    "### Model Configuration\n",
    "\n",
    "As with other demos we leverage the `context.Context` object to store all model and training parameters. \n",
    "One can set specific parameters using the `-set` command line flag.\n",
    "\n",
    "The [`imdb.CreateDefaultContext()`](https://github.com/gomlx/gomlx/blob/main/examples/imdb/train.go) method sets all the default values for the hyperparameters that may be used by any of the 3 model types. The parameter \"model\" specify the model type."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ad068f5e-46d9-411b-823c-0eae905da41c",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model types: [\"bow\" \"cnn\" \"transformer\"]\n",
      "\t\"/activation\": (string) \n",
      "\t\"/adam_dtype\": (string) \n",
      "\t\"/adam_epsilon\": (float64) 1e-07\n",
      "\t\"/batch_size\": (int) 32\n",
      "\t\"/cnn_dropout_rate\": (float64) 0.5\n",
      "\t\"/cnn_normalization\": (string) \n",
      "\t\"/cnn_num_layers\": (float64) 5\n",
      "\t\"/cosine_schedule_steps\": (int) 0\n",
      "\t\"/dropout_rate\": (float64) 0.1\n",
      "\t\"/eval_batch_size\": (int) 200\n",
      "\t\"/fnn_dropout_rate\": (float64) 0.3\n",
      "\t\"/fnn_normalization\": (string) \n",
      "\t\"/fnn_num_hidden_layers\": (int) 2\n",
      "\t\"/fnn_num_hidden_nodes\": (int) 32\n",
      "\t\"/fnn_residual\": (bool) true\n",
      "\t\"/imdb_content_max_len\": (int) 200\n",
      "\t\"/imdb_include_separators\": (bool) false\n",
      "\t\"/imdb_mask_word_task_weight\": (float64) 0\n",
      "\t\"/imdb_max_vocab\": (int) 20000\n",
      "\t\"/imdb_token_embedding_size\": (int) 32\n",
      "\t\"/imdb_use_unsupervised\": (bool) false\n",
      "\t\"/imdb_word_dropout_rate\": (float64) 0\n",
      "\t\"/l1_regularization\": (float64) 0\n",
      "\t\"/l2_regularization\": (float64) 0\n",
      "\t\"/learning_rate\": (float64) 0.0001\n",
      "\t\"/model\": (string) cnn\n",
      "\t\"/normalization\": (string) layer\n",
      "\t\"/num_checkpoints\": (int) 3\n",
      "\t\"/optimizer\": (string) adamw\n",
      "\t\"/plots\": (bool) true\n",
      "\t\"/train_steps\": (int) 5000\n",
      "\t\"/transformer_att_key_size\": (int) 8\n",
      "\t\"/transformer_dropout_rate\": (float64) -1\n",
      "\t\"/transformer_max_att_len\": (int) 200\n",
      "\t\"/transformer_num_att_heads\": (int) 2\n",
      "\t\"/transformer_num_att_layers\": (int) 1\n"
     ]
    }
   ],
   "source": [
    "import (\n",
    "    \"golang.org/x/exp/maps\"\n",
    "    \"github.com/gomlx/gomlx/pkg/ml/context\"\n",
    ")\n",
    "\n",
    "// settings is bound to a \"-set\" flag to be used to set context hyperparameters.\n",
    "var settings = commandline.CreateContextSettingsFlag(imdb.CreateDefaultContext(), \"set\")\n",
    "\n",
    "// ContextFromSettings is the default context (createDefaultContext) changed by -set flag.\n",
    "// It also returns the list of parameters changed by -set in paramsSet: we use this later to avoid loading over the values from checkpoints.\n",
    "func ContextFromSettings() (ctx *context.Context, paramsSet []string) {\n",
    "    ctx = imdb.CreateDefaultContext()\n",
    "    paramsSet = must.M1(commandline.ParseContextSettings(ctx, *settings))\n",
    "    return ctx, paramsSet\n",
    "}\n",
    "\n",
    "%% -set=\"model=cnn\"\n",
    "fmt.Printf(\"Model types: %q\\n\", maps.Keys(imdb.ValidModels))\n",
    "ctx, _ := ContextFromSettings()\n",
    "fmt.Println(commandline.SprintContextSettings(ctx))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b9eb816-6caf-46ea-805c-bd375834c55b",
   "metadata": {},
   "source": [
    "### Bag Of Words Model (bow)\n",
    "\n",
    "This is the simplest model we are going to train: it embeds each token of the sentence (default size of the is 32 numbers) sum them up, and pass that through a FNN.\n",
    "\n",
    "The [code in `imdb.BagOfWordsModelGraph`](https://github.com/gomlx/gomlx/blob/main/examples/imdb/model_bagofwords.go) looks like this:\n",
    "\n",
    "```go\n",
    "// BagOfWordsModelGraph builds the computation graph for the \"bag of words\" model: simply the sum of the embeddings\n",
    "// for each token included.\n",
    "func BagOfWordsModelGraph(ctx *context.Context, spec any, inputs []*Node) []*Node {\n",
    "\tembed, _ := EmbedTokensGraph(ctx, inputs[0])\n",
    "\n",
    "\t// Take the max over the content length, and put an FNN on top.\n",
    "\t// Shape transformation: [batch_size, content_len, embed_size] -> [batch_size, embed_size]\n",
    "\tembed = ReduceMax(embed, 1)\n",
    "\tlogits := fnn.New(ctx, embed, 1).Done()\n",
    "\treturn []*Node{logits}\n",
    "}\n",
    "```\n",
    "\n",
    "We played a bit with the hyperparameters to get to ~85% accuracy on the validation data.\n",
    "\n",
    "The [code for `imdb.TrainModel` is here](https://github.com/gomlx/gomlx/blob/main/examples/imdb/train.go).\n",
    "It's a straight forward GoMLX training loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "50cd517a-f592-4319-a9fd-acb8ab78124c",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> Loading previously generated preprocessed binary file.\n",
      "Loaded data from \"aclImdb.bin\": 100000 examples, 141088 unique tokens, 23727054 tokens in total.\n",
      "Backend \"stablehlo\":\tstablehlo:cuda - PJRT \"cuda\" plugin (/home/janpf/.local/lib/gomlx/pjrt/pjrt_c_api_cuda_plugin.so) v0.76 [StableHLO]\n",
      "Model: bow\n"
     ]
    },
    {
     "data": {
      "text/html": []
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "      \u001b[1m   7% [=>......................................] (189 steps/s) [9s:48s]\u001b[0m [step=719] [loss+=0.801] [~loss+=0.787] [~loss=0.676] [~acc=59.05%]         "
     ]
    },
    {
     "data": {
      "text/html": []
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "      \u001b[1m 100% [========================================] (130 steps/s)\u001b[0m [step=9999] [loss+=0.289] [~loss+=0.328] [~loss=0.277] [~acc=88.51%]        1%]        \n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<p><b>Metric: accuracy</b></p>\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div id=\"9d2ecc7e\"></div>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<script charset=\"UTF-8\">\n",
       "(() => {\n",
       "\tconst src=\"https://cdn.plot.ly/plotly-2.34.0.min.js\";\n",
       "\tvar runJSFn = function(module) {\n",
       "\t\t\n",
       "\tif (!module) {\n",
       "\t\tmodule = window.Plotly;\n",
       "\t}\n",
       "\tlet data = JSON.parse('{\"data\":[{\"type\":\"scatter\",\"line\":{\"shape\":\"linear\"},\"mode\":\"lines+markers\",\"name\":\"Train: Moving Average Accuracy\",\"x\":[201,441,729,1075,1490,1988,2586,3304,4166,5200,6441,7930,9717,10000],\"y\":[0.5358490943908691,0.5523581504821777,0.5917324423789978,0.6655932664871216,0.7119168043136597,0.7508869171142578,0.7796301245689392,0.8008884191513062,0.8283624053001404,0.8380887508392334,0.8598978519439697,0.8700001835823059,0.8859584927558899,0.8851222991943359]},{\"type\":\"scatter\",\"line\":{\"shape\":\"linear\"},\"mode\":\"lines+markers\",\"name\":\"Mean Accuracy on train-eval\",\"x\":[201,441,729,1075,1490,1988,2586,3304,4166,5200,6441,7930,9717,10000],\"y\":[0.5,0.6023199558258057,0.7239599823951721,0.7781199812889099,0.7988799810409546,0.8180800080299377,0.8388400077819824,0.855679988861084,0.870199978351593,0.8874399662017822,0.9030399918556213,0.9175999760627747,0.9322800040245056,0.9347999691963196]},{\"type\":\"scatter\",\"line\":{\"shape\":\"linear\"},\"mode\":\"lines+markers\",\"name\":\"Mean Accuracy on test-eval\",\"x\":[201,441,729,1075,1490,1988,2586,3304,4166,5200,6441,7930,9717,10000],\"y\":[0.5,0.5956000089645386,0.6994799971580505,0.75764000415802,0.7768799662590027,0.7945599555969238,0.814079999923706,0.8258399963378906,0.8337599635124207,0.8410799503326416,0.848360002040863,0.8532800078392029,0.8562399744987488,0.8571599721908569]}],\"layout\":{\"legend\":{},\"title\":{\"text\":\"accuracy\"},\"xaxis\":{\"showgrid\":true,\"title\":{\"text\":\"Steps\"},\"type\":\"log\"},\"yaxis\":{\"showgrid\":true,\"type\":\"log\"}}}');\n",
       "\tmodule.newPlot('9d2ecc7e', data);\n",
       "\n",
       "\t}\n",
       "\t\n",
       "    if (typeof requirejs === \"function\") {\n",
       "        // Use RequireJS to load module.\n",
       "\t\tlet srcWithoutExtension = src.substring(0, src.lastIndexOf(\".js\"));\n",
       "        requirejs.config({\n",
       "            paths: {\n",
       "                'plotly': srcWithoutExtension\n",
       "            }\n",
       "        });\n",
       "        require(['plotly'], function(plotly) {\n",
       "            runJSFn(plotly)\n",
       "        });\n",
       "        return\n",
       "    }\n",
       "\n",
       "\tvar currentScripts = document.head.getElementsByTagName(\"script\");\n",
       "\tfor (const idx in currentScripts) {\n",
       "\t\tlet script = currentScripts[idx];\n",
       "\t\tif (script.src == src) {\n",
       "\t\t\trunJSFn(null);\n",
       "\t\t\treturn;\n",
       "\t\t}\n",
       "\t}\n",
       "\n",
       "\tvar script = document.createElement(\"script\");\n",
       "\n",
       "\tscript.charset = \"utf-8\";\n",
       "\t\n",
       "\tscript.src = src;\n",
       "\tscript.onload = script.onreadystatechange = function () { runJSFn(null); };\n",
       "\tdocument.head.appendChild(script);\t\n",
       "})();\n",
       "</script>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<p><b>Metric: loss</b></p>\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div id=\"24c25ef9\"></div>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<script charset=\"UTF-8\">\n",
       "(() => {\n",
       "\tconst src=\"https://cdn.plot.ly/plotly-2.34.0.min.js\";\n",
       "\tvar runJSFn = function(module) {\n",
       "\t\t\n",
       "\tif (!module) {\n",
       "\t\tmodule = window.Plotly;\n",
       "\t}\n",
       "\tlet data = JSON.parse('{\"data\":[{\"type\":\"scatter\",\"line\":{\"shape\":\"linear\"},\"mode\":\"lines+markers\",\"name\":\"Train: Batch Loss+Regularization\",\"x\":[201,441,729,1075,1490,1988,2586,3304,4166,5200,6441,7930,9717,10000],\"y\":[0.8091092109680176,0.7930521965026855,0.7831490635871887,0.7353273630142212,0.6770973801612854,0.5083365440368652,0.588036060333252,0.5385592579841614,0.4370734691619873,0.46654200553894043,0.40635690093040466,0.3392995595932007,0.2635926306247711,0.288725346326828]},{\"type\":\"scatter\",\"line\":{\"shape\":\"linear\"},\"mode\":\"lines+markers\",\"name\":\"Train: Moving Average Loss+Regularization\",\"x\":[201,441,729,1075,1490,1988,2586,3304,4166,5200,6441,7930,9717,10000],\"y\":[0.8197277784347534,0.8078547716140747,0.7862991690635681,0.7498175501823425,0.6896934509277344,0.6199910044670105,0.5682920813560486,0.5200974941253662,0.4713091552257538,0.4470917582511902,0.3953598737716675,0.3706754744052887,0.32847920060157776,0.3283360004425049]},{\"type\":\"scatter\",\"line\":{\"shape\":\"linear\"},\"mode\":\"lines+markers\",\"name\":\"Train: Moving Average Loss\",\"x\":[201,441,729,1075,1490,1988,2586,3304,4166,5200,6441,7930,9717,10000],\"y\":[0.6899662017822266,0.6867415308952332,0.67488694190979,0.64781653881073,0.5953993797302246,0.5317043662071228,0.4856448769569397,0.44351547956466675,0.4004043936729431,0.3815089166164398,0.3353327810764313,0.3150138854980469,0.2769995629787445,0.277442991733551]},{\"type\":\"scatter\",\"line\":{\"shape\":\"linear\"},\"mode\":\"lines+markers\",\"name\":\"Mean Loss+Regularization on train-eval\",\"x\":[201,441,729,1075,1490,1988,2586,3304,4166,5200,6441,7930,9717,10000],\"y\":[0.8173539638519287,0.8000450730323792,0.7749971747398376,0.7311708331108093,0.6589702367782593,0.5772620439529419,0.5071929693222046,0.4543220102787018,0.40817660093307495,0.36640605330467224,0.32839053869247437,0.28908172249794006,0.250699907541275,0.24582359194755554]},{\"type\":\"scatter\",\"line\":{\"shape\":\"linear\"},\"mode\":\"lines+markers\",\"name\":\"Mean Loss on train-eval\",\"x\":[201,441,729,1075,1490,1988,2586,3304,4166,5200,6441,7930,9717,10000],\"y\":[0.6909013390541077,0.6825214624404907,0.6667217016220093,0.631476640701294,0.5661256313323975,0.4900237023830414,0.42544084787368774,0.37850528955459595,0.33785006403923035,0.3013964891433716,0.26860544085502625,0.2337370216846466,0.19934460520744324,0.1951548159122467]},{\"type\":\"scatter\",\"line\":{\"shape\":\"linear\"},\"mode\":\"lines+markers\",\"name\":\"Mean Loss+Regularization on test-eval\",\"x\":[201,441,729,1075,1490,1988,2586,3304,4166,5200,6441,7930,9717,10000],\"y\":[0.817565381526947,0.8011260032653809,0.7775143384933472,0.7362269163131714,0.6686755418777466,0.594642162322998,0.5331944227218628,0.4905681610107422,0.45685794949531555,0.4309251308441162,0.4109187722206116,0.3940573036670685,0.3821627199649811,0.38088271021842957]},{\"type\":\"scatter\",\"line\":{\"shape\":\"linear\"},\"mode\":\"lines+markers\",\"name\":\"Mean Loss on test-eval\",\"x\":[201,441,729,1075,1490,1988,2586,3304,4166,5200,6441,7930,9717,10000],\"y\":[0.6911131739616394,0.6836026310920715,0.6692390441894531,0.6365325450897217,0.5758306980133057,0.5074043869972229,0.4514424204826355,0.4147513210773468,0.3865320086479187,0.3659161627292633,0.3511337339878082,0.33871224522590637,0.3308075964450836,0.3302136957645416]}],\"layout\":{\"legend\":{},\"title\":{\"text\":\"loss\"},\"xaxis\":{\"showgrid\":true,\"title\":{\"text\":\"Steps\"},\"type\":\"log\"},\"yaxis\":{\"showgrid\":true,\"type\":\"log\"}}}');\n",
       "\tmodule.newPlot('24c25ef9', data);\n",
       "\n",
       "\t}\n",
       "\t\n",
       "    if (typeof requirejs === \"function\") {\n",
       "        // Use RequireJS to load module.\n",
       "\t\tlet srcWithoutExtension = src.substring(0, src.lastIndexOf(\".js\"));\n",
       "        requirejs.config({\n",
       "            paths: {\n",
       "                'plotly': srcWithoutExtension\n",
       "            }\n",
       "        });\n",
       "        require(['plotly'], function(plotly) {\n",
       "            runJSFn(plotly)\n",
       "        });\n",
       "        return\n",
       "    }\n",
       "\n",
       "\tvar currentScripts = document.head.getElementsByTagName(\"script\");\n",
       "\tfor (const idx in currentScripts) {\n",
       "\t\tlet script = currentScripts[idx];\n",
       "\t\tif (script.src == src) {\n",
       "\t\t\trunJSFn(null);\n",
       "\t\t\treturn;\n",
       "\t\t}\n",
       "\t}\n",
       "\n",
       "\tvar script = document.createElement(\"script\");\n",
       "\n",
       "\tscript.charset = \"utf-8\";\n",
       "\t\n",
       "\tscript.src = src;\n",
       "\tscript.onload = script.onreadystatechange = function () { runJSFn(null); };\n",
       "\tdocument.head.appendChild(script);\t\n",
       "})();\n",
       "</script>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t[Step 10000] median train step: 4563 microseconds\n",
      "\n",
      "Results on train-eval:\n",
      "\tMean Loss+Regularization (#loss+): 0.246\n",
      "\tMean Loss (#loss): 0.195\n",
      "\tMean Accuracy (#acc): 93.48%\n",
      "Results on test-eval:\n",
      "\tMean Loss+Regularization (#loss+): 0.381\n",
      "\tMean Loss (#loss): 0.33\n",
      "\tMean Accuracy (#acc): 85.72%\n"
     ]
    }
   ],
   "source": [
    "%% --set=\"model=bow;l2_regularization=1e-3;learning_rate=1e-4;normalization=none;train_steps=10000\"\n",
    "ctx, paramsSet := ContextFromSettings()\n",
    "imdb.TrainModel(ctx, *flagDataDir, *flagCheckpoint, paramsSet, *flagEval, *flagVerbosity)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62e0d8ed-31ed-49ba-8b6f-1518d3df8a8d",
   "metadata": {},
   "source": [
    "### Convolution Model (cnn)\n",
    "\n",
    "The function [`imdb.CnnModelGraph`](https://github.com/gomlx/gomlx/blob/main/examples/imdb/model_cnn.go) creates a 1D convolution model, with arbitrary number of convolutions. After the convolution, it behaves the same way as the Bag Of Words model.\n",
    "\n",
    "The core of the convolution model looks like this:\n",
    "\n",
    "```go\n",
    "\t// 1D Convolution: embed is [batch_size, content_len, embed_size].\n",
    "\tnumConvolutions := context.GetParamOr(ctx, \"cnn_num_layers\", 5)\n",
    "\tlogits := embed\n",
    "\tfor convIdx := range numConvolutions {\n",
    "\t\tctx := ctx.Inf(\"%03d_conv\", convIdx)\n",
    "\t\tresidual := logits\n",
    "\t\tif convIdx > 0 {\n",
    "\t\t\tlogits = NormalizeSequence(ctx, logits)\n",
    "\t\t}\n",
    "\t\tlogits = layers.Convolution(ctx, embed).KernelSize(7).Filters(embedSize).Strides(1).Done()\n",
    "\t\tlogits = activations.ApplyFromContext(ctx, logits)\n",
    "\t\tif dropoutNode != nil {\n",
    "\t\t\tlogits = layers.Dropout(ctx, logits, dropoutNode)\n",
    "\t\t}\n",
    "\t\tif residual.Shape().Equal(logits.Shape()) {\n",
    "\t\t\tlogits = Add(logits, residual)\n",
    "\t\t}\n",
    "\t}\n",
    "\n",
    "\t// Take the max over the content length, and put an FNN on top.\n",
    "\t// Shape transformation: [batch_size, content_len, embed_size] -> [batch_size, embed_size]\n",
    "\tlogits = ReduceMax(logits, 1)\n",
    "\tlogits = fnn.New(ctx, logits, 1).Done()\n",
    "\tlogits.AssertDims(batchSize, 1)\n",
    "```\n",
    "\n",
    "Notice how well it can overfit to the training data ... but it doesn't help the test results. To improve this one needs some careful regularization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f77a72a8-2872-4c56-9065-e245deb85169",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> Loading previously generated preprocessed binary file.\n",
      "Loaded data from \"aclImdb.bin\": 100000 examples, 141088 unique tokens, 23727054 tokens in total.\n",
      "Backend \"stablehlo\":\tstablehlo:cuda - PJRT \"cuda\" plugin (/home/janpf/.local/lib/gomlx/pjrt/pjrt_c_api_cuda_plugin.so) v0.76 [StableHLO]\n",
      "Model: cnn\n"
     ]
    },
    {
     "data": {
      "text/html": []
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "      \u001b[1m   7% [=>......................................] (165 steps/s) [14s:56s]\u001b[0m [step=719] [loss+=1.51] [~loss+=1.46] [~loss=0.781] [~acc=51.26%]          "
     ]
    },
    {
     "data": {
      "text/html": []
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "      \u001b[1m 100% [========================================] (95 steps/s)\u001b[0m [step=9999] [loss+=0.278] [~loss+=0.306] [~loss=0.105] [~acc=96.54%]        47%]         \n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<p><b>Metric: accuracy</b></p>\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div id=\"a15238c2\"></div>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<script charset=\"UTF-8\">\n",
       "(() => {\n",
       "\tconst src=\"https://cdn.plot.ly/plotly-2.34.0.min.js\";\n",
       "\tvar runJSFn = function(module) {\n",
       "\t\t\n",
       "\tif (!module) {\n",
       "\t\tmodule = window.Plotly;\n",
       "\t}\n",
       "\tlet data = JSON.parse('{\"data\":[{\"type\":\"scatter\",\"line\":{\"shape\":\"linear\"},\"mode\":\"lines+markers\",\"name\":\"Train: Moving Average Accuracy\",\"x\":[201,441,729,1075,1490,1988,2586,3304,4166,5200,6441,7930,9717,10000],\"y\":[0.5038279891014099,0.499844491481781,0.5129234194755554,0.5189930200576782,0.5507900714874268,0.627107560634613,0.7622535824775696,0.8339856266975403,0.8785673379898071,0.9014359712600708,0.930389404296875,0.9524072408676147,0.969953179359436,0.9654009938240051]},{\"type\":\"scatter\",\"line\":{\"shape\":\"linear\"},\"mode\":\"lines+markers\",\"name\":\"Mean Accuracy on train-eval\",\"x\":[201,441,729,1075,1490,1988,2586,3304,4166,5200,6441,7930,9717,10000],\"y\":[0.5,0.5140799880027771,0.5026400089263916,0.5169999599456787,0.5942400097846985,0.7672399878501892,0.8382399678230286,0.8890799880027771,0.9214400053024292,0.9450799822807312,0.9648399949073792,0.9788399934768677,0.9861999750137329,0.9886799454689026]},{\"type\":\"scatter\",\"line\":{\"shape\":\"linear\"},\"mode\":\"lines+markers\",\"name\":\"Mean Accuracy on test-eval\",\"x\":[201,441,729,1075,1490,1988,2586,3304,4166,5200,6441,7930,9717,10000],\"y\":[0.5,0.5107600092887878,0.502240002155304,0.5149199962615967,0.5682799816131592,0.7192800045013428,0.7915999889373779,0.8287599682807922,0.8434000015258789,0.8470399975776672,0.8453599810600281,0.8420000076293945,0.8331999778747559,0.8369999527931213]}],\"layout\":{\"legend\":{},\"title\":{\"text\":\"accuracy\"},\"xaxis\":{\"showgrid\":true,\"title\":{\"text\":\"Steps\"},\"type\":\"log\"},\"yaxis\":{\"showgrid\":true,\"type\":\"log\"}}}');\n",
       "\tmodule.newPlot('a15238c2', data);\n",
       "\n",
       "\t}\n",
       "\t\n",
       "    if (typeof requirejs === \"function\") {\n",
       "        // Use RequireJS to load module.\n",
       "\t\tlet srcWithoutExtension = src.substring(0, src.lastIndexOf(\".js\"));\n",
       "        requirejs.config({\n",
       "            paths: {\n",
       "                'plotly': srcWithoutExtension\n",
       "            }\n",
       "        });\n",
       "        require(['plotly'], function(plotly) {\n",
       "            runJSFn(plotly)\n",
       "        });\n",
       "        return\n",
       "    }\n",
       "\n",
       "\tvar currentScripts = document.head.getElementsByTagName(\"script\");\n",
       "\tfor (const idx in currentScripts) {\n",
       "\t\tlet script = currentScripts[idx];\n",
       "\t\tif (script.src == src) {\n",
       "\t\t\trunJSFn(null);\n",
       "\t\t\treturn;\n",
       "\t\t}\n",
       "\t}\n",
       "\n",
       "\tvar script = document.createElement(\"script\");\n",
       "\n",
       "\tscript.charset = \"utf-8\";\n",
       "\t\n",
       "\tscript.src = src;\n",
       "\tscript.onload = script.onreadystatechange = function () { runJSFn(null); };\n",
       "\tdocument.head.appendChild(script);\t\n",
       "})();\n",
       "</script>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<p><b>Metric: loss</b></p>\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div id=\"8149653c\"></div>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<script charset=\"UTF-8\">\n",
       "(() => {\n",
       "\tconst src=\"https://cdn.plot.ly/plotly-2.34.0.min.js\";\n",
       "\tvar runJSFn = function(module) {\n",
       "\t\t\n",
       "\tif (!module) {\n",
       "\t\tmodule = window.Plotly;\n",
       "\t}\n",
       "\tlet data = JSON.parse('{\"data\":[{\"type\":\"scatter\",\"line\":{\"shape\":\"linear\"},\"mode\":\"lines+markers\",\"name\":\"Train: Batch Loss+Regularization\",\"x\":[201,441,729,1075,1490,1988,2586,3304,4166,5200,6441,7930,9717,10000],\"y\":[1.4257235527038574,1.4992945194244385,1.3979971408843994,1.4015116691589355,1.201706886291504,1.102069616317749,0.8799208402633667,0.6195549964904785,0.7266563177108765,0.5455760955810547,0.4851660132408142,0.5179839134216309,0.33373770117759705,0.27759772539138794]},{\"type\":\"scatter\",\"line\":{\"shape\":\"linear\"},\"mode\":\"lines+markers\",\"name\":\"Train: Moving Average Loss+Regularization\",\"x\":[201,441,729,1075,1490,1988,2586,3304,4166,5200,6441,7930,9717,10000],\"y\":[1.6190334558486938,1.5363991260528564,1.453014612197876,1.361562967300415,1.2669110298156738,1.126526117324829,0.9137341976165771,0.7672454714775085,0.6455115079879761,0.5621262192726135,0.45921221375465393,0.36833661794662476,0.29661041498184204,0.3058527708053589]},{\"type\":\"scatter\",\"line\":{\"shape\":\"linear\"},\"mode\":\"lines+markers\",\"name\":\"Train: Moving Average Loss\",\"x\":[201,441,729,1075,1490,1988,2586,3304,4166,5200,6441,7930,9717,10000],\"y\":[0.8677921295166016,0.8179896473884583,0.7795875072479248,0.7438068389892578,0.7139736413955688,0.6434528231620789,0.4932347536087036,0.3920222520828247,0.3073444664478302,0.25872939825057983,0.19032002985477448,0.13289333879947662,0.09162840992212296,0.10491303354501724]},{\"type\":\"scatter\",\"line\":{\"shape\":\"linear\"},\"mode\":\"lines+markers\",\"name\":\"Mean Loss+Regularization on train-eval\",\"x\":[201,441,729,1075,1490,1988,2586,3304,4166,5200,6441,7930,9717,10000],\"y\":[1.4989075660705566,1.3871233463287354,1.3403433561325073,1.2720565795898438,1.1770614385604858,0.9644309282302856,0.7957513332366943,0.6733590960502625,0.5612586140632629,0.46329671144485474,0.37716013193130493,0.3047415018081665,0.2503986060619354,0.23935885727405548]},{\"type\":\"scatter\",\"line\":{\"shape\":\"linear\"},\"mode\":\"lines+markers\",\"name\":\"Mean Loss on train-eval\",\"x\":[201,441,729,1075,1490,1988,2586,3304,4166,5200,6441,7930,9717,10000],\"y\":[0.759210467338562,0.6836805939674377,0.6829113960266113,0.6703471541404724,0.6392611861228943,0.49400317668914795,0.3833142817020416,0.3032577335834503,0.22689561545848846,0.1629633903503418,0.11075165122747421,0.07122274488210678,0.0468786358833313,0.039822012186050415]},{\"type\":\"scatter\",\"line\":{\"shape\":\"linear\"},\"mode\":\"lines+markers\",\"name\":\"Mean Loss+Regularization on test-eval\",\"x\":[201,441,729,1075,1490,1988,2586,3304,4166,5200,6441,7930,9717,10000],\"y\":[1.5027552843093872,1.394626259803772,1.3500741720199585,1.286898136138916,1.2054553031921387,1.0149005651474,0.8846902251243591,0.8378227353096008,0.7977476119995117,0.7619814872741699,0.7720106840133667,0.8334422707557678,0.953571081161499,0.9424947500228882]},{\"type\":\"scatter\",\"line\":{\"shape\":\"linear\"},\"mode\":\"lines+markers\",\"name\":\"Mean Loss on test-eval\",\"x\":[201,441,729,1075,1490,1988,2586,3304,4166,5200,6441,7930,9717,10000],\"y\":[0.7630581259727478,0.6911835074424744,0.69264155626297,0.6851896643638611,0.6676549911499023,0.5444731116294861,0.4722535014152527,0.46772146224975586,0.46338513493537903,0.4616492986679077,0.5056023001670837,0.5999233722686768,0.7500501871109009,0.7429568767547607]}],\"layout\":{\"legend\":{},\"title\":{\"text\":\"loss\"},\"xaxis\":{\"showgrid\":true,\"title\":{\"text\":\"Steps\"},\"type\":\"log\"},\"yaxis\":{\"showgrid\":true,\"type\":\"log\"}}}');\n",
       "\tmodule.newPlot('8149653c', data);\n",
       "\n",
       "\t}\n",
       "\t\n",
       "    if (typeof requirejs === \"function\") {\n",
       "        // Use RequireJS to load module.\n",
       "\t\tlet srcWithoutExtension = src.substring(0, src.lastIndexOf(\".js\"));\n",
       "        requirejs.config({\n",
       "            paths: {\n",
       "                'plotly': srcWithoutExtension\n",
       "            }\n",
       "        });\n",
       "        require(['plotly'], function(plotly) {\n",
       "            runJSFn(plotly)\n",
       "        });\n",
       "        return\n",
       "    }\n",
       "\n",
       "\tvar currentScripts = document.head.getElementsByTagName(\"script\");\n",
       "\tfor (const idx in currentScripts) {\n",
       "\t\tlet script = currentScripts[idx];\n",
       "\t\tif (script.src == src) {\n",
       "\t\t\trunJSFn(null);\n",
       "\t\t\treturn;\n",
       "\t\t}\n",
       "\t}\n",
       "\n",
       "\tvar script = document.createElement(\"script\");\n",
       "\n",
       "\tscript.charset = \"utf-8\";\n",
       "\t\n",
       "\tscript.src = src;\n",
       "\tscript.onload = script.onreadystatechange = function () { runJSFn(null); };\n",
       "\tdocument.head.appendChild(script);\t\n",
       "})();\n",
       "</script>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t[Step 10000] median train step: 4881 microseconds\n",
      "\n",
      "Results on train-eval:\n",
      "\tMean Loss+Regularization (#loss+): 0.239\n",
      "\tMean Loss (#loss): 0.0398\n",
      "\tMean Accuracy (#acc): 98.87%\n",
      "Results on test-eval:\n",
      "\tMean Loss+Regularization (#loss+): 0.942\n",
      "\tMean Loss (#loss): 0.743\n",
      "\tMean Accuracy (#acc): 83.70%\n"
     ]
    }
   ],
   "source": [
    "%% --set=\"model=cnn;l2_regularization=1e-3;learning_rate=1e-4;normalization=layer;train_steps=10000\"\n",
    "ctx, paramsSet := ContextFromSettings()\n",
    "imdb.TrainModel(ctx, *flagDataDir, *flagCheckpoint, paramsSet, *flagEval, *flagVerbosity)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9adcb51-2190-4cfc-9224-07f7331bae41",
   "metadata": {},
   "source": [
    "### Transformer Model\n",
    "\n",
    "Finally a Transformer version of the model, as defined in the [\"Attention Is All You Need\"](https://arxiv.org/abs/1706.03762) famous paper. \n",
    "\n",
    "Notice it's not significantly better than our previous simple Bag-Of-Words model. Likely because there is not enough data for the transformer to make any difference. The success of transformers in large-language-models is in large part due to the training with huge amounts of unsupervised (or self-supervised) data, but that is beyond the scope of this small test.\n",
    "\n",
    "The code is in [`imdb.TransformerModelGraph`](https://github.com/gomlx/gomlx/blob/main/examples/imdb/model_transformer.go), and the core of it looks like this:\n",
    "\n",
    "```go\n",
    "    ...\n",
    "\t// Add the requested number of attention layers.\n",
    "\tnumAttLayers := context.GetParamOr(ctx, \"transformer_num_att_layers\", 1)\n",
    "\tnumAttHeads := context.GetParamOr(ctx, \"transformer_num_att_heads\", 2)\n",
    "\tattKeySize := context.GetParamOr(ctx, \"transformer_att_key_size\", 8)\n",
    "\tfor layerNum := range numAttLayers {\n",
    "\t\t// Each layer in its own scope.\n",
    "\t\tctx := ctx.Inf(\"%03d_attention_layer\", layerNum)\n",
    "\t\tresidual := embed\n",
    "\t\tembed = layers.MultiHeadAttention(ctx.In(\"000_attention\"), embed, embed, embed, numAttHeads, attKeySize).\n",
    "\t\t\tSetKeyMask(mask).SetQueryMask(mask).\n",
    "\t\t\tSetOutputDim(embedSize).\n",
    "\t\t\tSetValueHeadDim(embedSize).Done()\n",
    "\t\tif dropoutNode != nil {\n",
    "\t\t\tembed = layers.Dropout(ctx.In(\"001_dropout\"), embed, dropoutNode)\n",
    "\t\t}\n",
    "\t\tembed = NormalizeSequence(ctx.In(\"002_normalization\"), embed)\n",
    "\t\tattentionOutput := embed\n",
    "\n",
    "\t\t// Transformers recipe: 2 dense layers after attention.\n",
    "\t\tembed = fnn.New(ctx.In(\"003_fnn\"), embed, embedSize).NumHiddenLayers(1, embedSize).Done()\n",
    "\t\tif dropoutNode != nil {\n",
    "\t\t\tembed = layers.Dropout(ctx.In(\"004_dropout\"), embed, dropoutNode)\n",
    "\t\t}\n",
    "\t\tembed = Add(embed, attentionOutput)\n",
    "\t\tembed = NormalizeSequence(ctx.In(\"005_normalization\"), embed)\n",
    "\n",
    "\t\t// Residual connection:\n",
    "\t\tif layerNum > 0 {\n",
    "\t\t\tembed = Add(residual, embed)\n",
    "\t\t}\n",
    "\t}\n",
    "    ...\n",
    "```\n",
    "\n",
    "With only 5000 steps we got ~87% on the test data -- and significant overfitting as well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "20ac0032-de61-4573-88db-fbac6730ac9f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "> Loading previously generated preprocessed binary file.\n",
      "Loaded data from \"aclImdb.bin\": 100000 examples, 141088 unique tokens, 23727054 tokens in total.\n",
      "Backend \"stablehlo\":\tstablehlo:cuda - PJRT \"cuda\" plugin (/home/janpf/.local/lib/gomlx/pjrt/pjrt_c_api_cuda_plugin.so) v0.76 [StableHLO]\n",
      "Model: transformer\n"
     ]
    },
    {
     "data": {
      "text/html": []
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "      \u001b[1m  14% [====>...................................] (161 steps/s) [13s:26s]\u001b[0m [step=724] [loss+=1.05] [~loss+=1.21] [~loss=0.816] [~acc=51.65%]        "
     ]
    },
    {
     "data": {
      "text/html": []
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "      \u001b[1m 100% [========================================] (82 steps/s)\u001b[0m [step=4999] [loss+=0.599] [~loss+=0.565] [~loss=0.227] [~acc=91.92%]        9%]        \n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<p><b>Metric: accuracy</b></p>\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div id=\"9d78bd5a\"></div>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<script charset=\"UTF-8\">\n",
       "(() => {\n",
       "\tconst src=\"https://cdn.plot.ly/plotly-2.34.0.min.js\";\n",
       "\tvar runJSFn = function(module) {\n",
       "\t\t\n",
       "\tif (!module) {\n",
       "\t\tmodule = window.Plotly;\n",
       "\t}\n",
       "\tlet data = JSON.parse('{\"data\":[{\"type\":\"scatter\",\"line\":{\"shape\":\"linear\"},\"mode\":\"lines+markers\",\"name\":\"Train: Moving Average Accuracy\",\"x\":[201,441,729,1075,1490,1988,2586,3304,4166,5000],\"y\":[0.5083816647529602,0.4997476637363434,0.5164817571640015,0.5482693910598755,0.6316553950309753,0.7418219447135925,0.8201721906661987,0.8732424378395081,0.9033474922180176,0.9192222952842712]},{\"type\":\"scatter\",\"line\":{\"shape\":\"linear\"},\"mode\":\"lines+markers\",\"name\":\"Mean Accuracy on train-eval\",\"x\":[201,441,729,1075,1490,1988,2586,3304,4166,5000],\"y\":[0.5,0.5,0.55867999792099,0.7508400082588196,0.7557199597358704,0.7917199730873108,0.8660399913787842,0.8703599572181702,0.8988399505615234,0.8812399506568909]},{\"type\":\"scatter\",\"line\":{\"shape\":\"linear\"},\"mode\":\"lines+markers\",\"name\":\"Mean Accuracy on test-eval\",\"x\":[201,441,729,1075,1490,1988,2586,3304,4166,5000],\"y\":[0.5,0.5,0.5613200068473816,0.6975599527359009,0.730239987373352,0.7607600092887878,0.8255199790000916,0.8307200074195862,0.8487599492073059,0.8269999623298645]}],\"layout\":{\"legend\":{},\"title\":{\"text\":\"accuracy\"},\"xaxis\":{\"showgrid\":true,\"title\":{\"text\":\"Steps\"},\"type\":\"log\"},\"yaxis\":{\"showgrid\":true,\"type\":\"log\"}}}');\n",
       "\tmodule.newPlot('9d78bd5a', data);\n",
       "\n",
       "\t}\n",
       "\t\n",
       "    if (typeof requirejs === \"function\") {\n",
       "        // Use RequireJS to load module.\n",
       "\t\tlet srcWithoutExtension = src.substring(0, src.lastIndexOf(\".js\"));\n",
       "        requirejs.config({\n",
       "            paths: {\n",
       "                'plotly': srcWithoutExtension\n",
       "            }\n",
       "        });\n",
       "        require(['plotly'], function(plotly) {\n",
       "            runJSFn(plotly)\n",
       "        });\n",
       "        return\n",
       "    }\n",
       "\n",
       "\tvar currentScripts = document.head.getElementsByTagName(\"script\");\n",
       "\tfor (const idx in currentScripts) {\n",
       "\t\tlet script = currentScripts[idx];\n",
       "\t\tif (script.src == src) {\n",
       "\t\t\trunJSFn(null);\n",
       "\t\t\treturn;\n",
       "\t\t}\n",
       "\t}\n",
       "\n",
       "\tvar script = document.createElement(\"script\");\n",
       "\n",
       "\tscript.charset = \"utf-8\";\n",
       "\t\n",
       "\tscript.src = src;\n",
       "\tscript.onload = script.onreadystatechange = function () { runJSFn(null); };\n",
       "\tdocument.head.appendChild(script);\t\n",
       "})();\n",
       "</script>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<p><b>Metric: loss</b></p>\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div id=\"75ba0d71\"></div>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<script charset=\"UTF-8\">\n",
       "(() => {\n",
       "\tconst src=\"https://cdn.plot.ly/plotly-2.34.0.min.js\";\n",
       "\tvar runJSFn = function(module) {\n",
       "\t\t\n",
       "\tif (!module) {\n",
       "\t\tmodule = window.Plotly;\n",
       "\t}\n",
       "\tlet data = JSON.parse('{\"data\":[{\"type\":\"scatter\",\"line\":{\"shape\":\"linear\"},\"mode\":\"lines+markers\",\"name\":\"Train: Batch Loss+Regularization\",\"x\":[201,441,729,1075,1490,1988,2586,3304,4166,5000],\"y\":[3.506824254989624,1.6666475534439087,1.0896365642547607,1.0731967687606812,1.035282850265503,0.9393253326416016,0.7424170970916748,0.7176443338394165,0.5886024236679077,0.5988227128982544]},{\"type\":\"scatter\",\"line\":{\"shape\":\"linear\"},\"mode\":\"lines+markers\",\"name\":\"Train: Moving Average Loss+Regularization\",\"x\":[201,441,729,1075,1490,1988,2586,3304,4166,5000],\"y\":[6.1751017570495605,2.4864602088928223,1.208648681640625,1.0786584615707397,1.0367670059204102,0.9707230925559998,0.8524488806724548,0.7137212157249451,0.6135207414627075,0.5645744204521179]},{\"type\":\"scatter\",\"line\":{\"shape\":\"linear\"},\"mode\":\"lines+markers\",\"name\":\"Train: Moving Average Loss\",\"x\":[201,441,729,1075,1490,1988,2586,3304,4166,5000],\"y\":[5.764137268066406,2.0819644927978516,0.8109996914863586,0.6893489956855774,0.6562697887420654,0.5987900495529175,0.48761916160583496,0.35507431626319885,0.26414915919303894,0.22720715403556824]},{\"type\":\"scatter\",\"line\":{\"shape\":\"linear\"},\"mode\":\"lines+markers\",\"name\":\"Mean Loss+Regularization on train-eval\",\"x\":[201,441,729,1075,1490,1988,2586,3304,4166,5000],\"y\":[2.256531000137329,1.149381160736084,1.0880976915359497,1.0593465566635132,1.029381275177002,0.9721778035163879,0.866571307182312,0.7469726204872131,0.6519964337348938,0.6346895694732666]},{\"type\":\"scatter\",\"line\":{\"shape\":\"linear\"},\"mode\":\"lines+markers\",\"name\":\"Mean Loss on train-eval\",\"x\":[201,441,729,1075,1490,1988,2586,3304,4166,5000],\"y\":[1.848422884941101,0.7473310828208923,0.6929046511650085,0.6723279356956482,0.6508281826972961,0.6017296314239502,0.5026690363883972,0.38919541239738464,0.3038582503795624,0.29895177483558655]},{\"type\":\"scatter\",\"line\":{\"shape\":\"linear\"},\"mode\":\"lines+markers\",\"name\":\"Mean Loss+Regularization on test-eval\",\"x\":[201,441,729,1075,1490,1988,2586,3304,4166,5000],\"y\":[2.4154250621795654,1.150821328163147,1.0842758417129517,1.0627048015594482,1.0350865125656128,0.9842621684074402,0.8900580406188965,0.7847957611083984,0.712668240070343,0.7170745730400085]},{\"type\":\"scatter\",\"line\":{\"shape\":\"linear\"},\"mode\":\"lines+markers\",\"name\":\"Mean Loss on test-eval\",\"x\":[201,441,729,1075,1490,1988,2586,3304,4166,5000],\"y\":[2.0073165893554688,0.7487710118293762,0.689084529876709,0.675686240196228,0.6565330028533936,0.6138137578964233,0.526156485080719,0.4270189702510834,0.36453017592430115,0.38133659958839417]}],\"layout\":{\"legend\":{},\"title\":{\"text\":\"loss\"},\"xaxis\":{\"showgrid\":true,\"title\":{\"text\":\"Steps\"},\"type\":\"log\"},\"yaxis\":{\"showgrid\":true,\"type\":\"log\"}}}');\n",
       "\tmodule.newPlot('75ba0d71', data);\n",
       "\n",
       "\t}\n",
       "\t\n",
       "    if (typeof requirejs === \"function\") {\n",
       "        // Use RequireJS to load module.\n",
       "\t\tlet srcWithoutExtension = src.substring(0, src.lastIndexOf(\".js\"));\n",
       "        requirejs.config({\n",
       "            paths: {\n",
       "                'plotly': srcWithoutExtension\n",
       "            }\n",
       "        });\n",
       "        require(['plotly'], function(plotly) {\n",
       "            runJSFn(plotly)\n",
       "        });\n",
       "        return\n",
       "    }\n",
       "\n",
       "\tvar currentScripts = document.head.getElementsByTagName(\"script\");\n",
       "\tfor (const idx in currentScripts) {\n",
       "\t\tlet script = currentScripts[idx];\n",
       "\t\tif (script.src == src) {\n",
       "\t\t\trunJSFn(null);\n",
       "\t\t\treturn;\n",
       "\t\t}\n",
       "\t}\n",
       "\n",
       "\tvar script = document.createElement(\"script\");\n",
       "\n",
       "\tscript.charset = \"utf-8\";\n",
       "\t\n",
       "\tscript.src = src;\n",
       "\tscript.onload = script.onreadystatechange = function () { runJSFn(null); };\n",
       "\tdocument.head.appendChild(script);\t\n",
       "})();\n",
       "</script>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t[Step 5000] median train step: 5063 microseconds\n",
      "\n",
      "Results on train-eval:\n",
      "\tMean Loss+Regularization (#loss+): 0.635\n",
      "\tMean Loss (#loss): 0.299\n",
      "\tMean Accuracy (#acc): 88.12%\n",
      "Results on test-eval:\n",
      "\tMean Loss+Regularization (#loss+): 0.717\n",
      "\tMean Loss (#loss): 0.381\n",
      "\tMean Accuracy (#acc): 82.70%\n"
     ]
    }
   ],
   "source": [
    "%% --set=\"model=transformer;normalization=none;activation=swish;l2_regularization=1e-3;cnn_dropout_rate=0.5;fnn_dropout_rate=0.3;learning_rate=1e-4;train_steps=5000\"\n",
    "ctx, paramsSet := ContextFromSettings()\n",
    "imdb.TrainModel(ctx, *flagDataDir, *flagCheckpoint, paramsSet, *flagEval, *flagVerbosity)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Go (gonb)",
   "language": "go",
   "name": "gonb"
  },
  "language_info": {
   "codemirror_mode": "",
   "file_extension": ".go",
   "mimetype": "text/x-go",
   "name": "go",
   "nbconvert_exporter": "",
   "pygments_lexer": "",
   "version": "go1.25.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
